package cn.cew.bytespy.transfer;

import cn.cew.bytespy.comm.AgentParam;
import cn.cew.bytespy.comm.Operator;
import cn.cew.bytespy.utils.BaseUtils;
import javassist.*;
import javassist.bytecode.CodeAttribute;
import javassist.bytecode.LocalVariableAttribute;
import javassist.bytecode.MethodInfo;

import java.io.File;
import java.io.FileOutputStream;;
import java.lang.reflect.Method;
import java.security.ProtectionDomain;
import java.text.MessageFormat;


/**
 * @author chenenwei
 * @version 1.0.0
 * @description TODO
 * @createTime 2022-10-15 14:47:00
 */

public class MethodCostTransformer extends BaseTransformer {

    public MethodCostTransformer(AgentParam agentParam) {
        this.agentParam = agentParam;
    }

    @Override
    public byte[] doTransform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws Exception {

        ClassPool classPool = ClassPool.getDefault();
        // 导入包
        if (BaseUtils.isNotEmpty(agentParam.getImportPackages())) {
            for (String packageName : agentParam.getImportPackages()) {
                classPool.importPackage(packageName);
            }
        }

        String classNameDot = BaseUtils.packageToDot(className);
        CtClass ctClass = null;
        try {
            ctClass = classPool.getCtClass(classNameDot);
        } catch (NotFoundException e) {
            e.printStackTrace();
            return null;
        }
        if (ctClass.isFrozen()) {
            ctClass.defrost();
        }
        CtMethod targetMethod = null;
        CtMethod[] declaredMethods = ctClass.getDeclaredMethods();
        for (CtMethod declaredMethod : declaredMethods) {
            // 根据签名获取目标方法
            String descriptor = declaredMethod.getMethodInfo().getDescriptor();
            if (descriptor.equals(agentParam.getSignature()) && declaredMethod.getName().equals(agentParam.getMethodName())) {
                targetMethod = declaredMethod;
                break;
            }
        }
        if (BaseUtils.isNotEmpty(agentParam.getOperators())) {
            System.out.println("开始添加代码");
            for (String operator : agentParam.getOperators()) {
                // 方法参数
                if (operator.equals(Operator.METHOD_PARAM.getValue())) {
                    printMethodParam(targetMethod);
                    // 全局变量
                } else if (operator.equals(Operator.METHOD_GLOBAL_VAR.getValue())) {
                    printGlobalVar(targetMethod);
                    // 局部变量
                } else if (operator.equals(Operator.METHOD_LOCAL_VAR.getValue())) {
                    // 方法耗时
                } else if (operator.equals(Operator.METHOD_COST.getValue())) {
                    printMethodCost(classPool, targetMethod);
                    // 方法返回值
                } else if (operator.equals(Operator.METHOD_RETURN.getValue())) {

                }
            }
        }
        byte[] bytes = ctClass.toBytecode();
        FileOutputStream out = new FileOutputStream(new File("first.class"));
        out.write(bytes);
        out.flush();
        out.close();
        return bytes;
    }

    @Override
    public byte[] makeClass(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws NotFoundException, CannotCompileException {
        ClassPool classPool = ClassPool.getDefault();
        String cName = BaseUtils.packageToDot(className);
        CtClass tempClass = classPool.getCtClass(cName + "AgentTest");
        tempClass.defrost();
        CtMethod oldMethod = tempClass.getDeclaredMethod(agentParam.getMethodName());
        tempClass.removeMethod(oldMethod);
        CtMethod method = CtMethod.make(agentParam.getMethodBody(), tempClass);
        tempClass.addMethod(method);
//        byte[] bytes = tempClass.toBytecode();
//        FileOutputStream out = new FileOutputStream(new File("agent.class"));
//        out.write(bytes);
//        out.flush();
//        out.close();
        return null;
    }



    /**
     * 打印方法耗时统计
     */
    private void printMethodCost (ClassPool classPool, CtMethod targetMethod) throws Exception {

        targetMethod.addLocalVariable("start", CtClass.longType);
        targetMethod.addLocalVariable("second", classPool.get("java.lang.String"));
        targetMethod.insertBefore("start = System.currentTimeMillis();");
        targetMethod.insertAfter("second =  (System.currentTimeMillis() - start) + \"ms\" ;");
        String msg = makeMethodMsg(targetMethod.getName(), "方法耗时", "second");
        String callMethod = makeSend(msg);
        targetMethod.insertAfter(callMethod);
    }

    /**
     * 打印全局变量
     */
    private void printGlobalVar (CtMethod targetMethod) throws Exception{

        if (BaseUtils.isNotEmpty(agentParam.getGlobalVars())) {
            StringBuilder code = new StringBuilder();
            for (String var : agentParam.getGlobalVars()) {
                String msg = makeVarMsg(targetMethod.getName(), var, "全局变量", var);
                code.append(makeSend(msg));
            }
            targetMethod.insertBefore(code.toString());
        }
    }

    /**
     * 打印方法参数
     */
    private void printMethodParam (CtMethod targetMethod) throws Exception{

        String[] allParameterName = getAllParameterName(targetMethod);
        String data = "";
        if (BaseUtils.isNotEmpty(allParameterName)) {
            for (int length = allParameterName.length; length > 0; length--) {
                String msg = makeMethodMsg(targetMethod.getName(), "方法参数", allParameterName[length - 1]);
                targetMethod.insertBefore(makeSend(msg));
            }
        }
    }

    /**
     * 获取方法参数名称
     * @param method 方法
     * @return
     * @throws NotFoundException
     */
    public static String[] getAllParameterName(CtMethod method) throws NotFoundException {

        MethodInfo methodInfo = method.getMethodInfo();
        CodeAttribute codeAttribute = methodInfo.getCodeAttribute();
        LocalVariableAttribute attr = (LocalVariableAttribute) codeAttribute.getAttribute(LocalVariableAttribute.tag);
        int pos = Modifier.isStatic(method.getModifiers()) ? 0 : 1;
        String[] paramNames = new String[method.getParameterTypes().length];
        for (int i = 0; i < paramNames.length; i++) {
            paramNames[i] = attr.variableName(i + pos);
        }
        return paramNames;
    }
}
